# 根据文件扩展名，使用不同的库
import rarfile, glob, os, zipfile
import multiprocessing as mp
from functools import partial, lru_cache
import sys, locale


class Unzip:
    def __init__(self, path):
        self.path = path
        self.temp_dir = "temp"
        self.unzip_dir = "unzip"
        os.makedirs(self.unzip_dir, exist_ok=True)
        os.makedirs(self.temp_dir, exist_ok=True)

    def by2str(self, by):
        code = locale.getpreferredencoding()
        for font in ['utf-8', code]:
            try:
                s = by.decode(font)
                return s
            except Exception as e:
                print(e)
                continue
        return None

    def evenly_pwds_chunk(self):
        l = []
        for i in glob.glob('./pwd/*'):
            with open(i, "rb") as f:
                by = f.read()
                s = self.by2str(by)
                if s is not None:
                    pwds = s.split("\n")
                    l = l + pwds
        ll = list(set(l))
        cpu = os.cpu_count() - 1
        total = len(ll)
        print(f"total pwds=={int(total / 10000)}万")
        n = int(total / cpu)
        print(f"every list=={int(n / 10000)}万")
        for i in range(0, len(ll), n):
            yield ll[i:i + n]

    # @lru_cache is not working in multiprocess
    def open_zip(self, path):
        if self.path.endswith('.zip'):
            return zipfile.ZipFile(path)
        if self.path.endswith('.rar'):
            return rarfile.RarFile(path)

    def extract(self, pwds):
        fp = self.open_zip(self.path)
        for pwd in pwds:
            pwd = pwd.rstrip()
            try:
                fp.extractall(path=self.temp_dir, pwd=pwd.encode())
                # 在multiprocess中打不出来，命令行运可以，所以在运行conf中emulate terminal out console,上一个temp  dir 大家都写，不会有东西，对了写到unzip dir
                print('Success! =================>' + pwd, flush=True)
                fp.extractall(path=self.unzip_dir, pwd=pwd.encode())
                fp.close()
                return pwd
            except:
                continue
        fp.close()
        return

    # 使用partial，list代表的入参必须在最后p.map(partial(self.extract, path=path), l)
    def run(self):
        pwds_list = self.evenly_pwds_chunk()
        with mp.Pool() as p:
            res = p.map(self.extract, pwds_list)
            print(res)


if __name__ == '__main__':
    u = Unzip("b.rar")
    u.run()
